import sys
import cv2
import os
import numpy as np
import pickle
from datetime import datetime
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import *


class FaceRecognitionSystem(QWidget):
    def __init__(self):
        super().__init__()
        self.cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        self.face_recognizer = cv2.face.LBPHFaceRecognizer_create()
        self.label_map = {}  # 原始结构保留
        self.frame_counter = 0
        self.is_paused = False
        self.conf_threshold = 80  # 默认置信度阈值

        self.init_ui()
        self.load_model()  # 启动时加载模型
        self.timer = self.startTimer(30)

    def init_ui(self):
        """原始界面代码完整保留（仅补充注释）"""
        self.setWindowTitle("智能人脸识别考勤系统")
        self.setGeometry(100, 100, 1024, 600)
        self.setStyleSheet("""
        QWidget { 
            background-image: url(img/bg.jpg); 
            background-repeat: no-repeat; 
            background-attachment: fixed; 
            background-size: cover; 
            font-family: "微软雅黑"; 
        }
        QLabel { color: #2c3e50; }
        QPushButton { 
            background-color: #3498db; 
            color: white; 
            border-radius: 5px; 
            padding: 8px 20px; 
            font-size: 14px; 
        }
         self.pushButton.setStyleSheet(QPushButton:hover{ 
            background-color: red; 
            color: white; 
            })
        QLineEdit { border: 2px solid #bdc3c7; border-radius: 5px; padding: 8px; font-size: 14px; }
        """)
        main_layout = QHBoxLayout()

        # 左侧视频显示区（原始代码完整保留）
        video_panel = QWidget()
        video_layout = QVBoxLayout()
        self.video_label = QLabel()
        self.video_label.setFixedSize(700, 500)
        self.video_label.setStyleSheet("background-color: white; border-radius: 10px;")
        video_layout.addWidget(self.video_label, alignment=Qt.AlignCenter)
        video_panel.setLayout(video_layout)
        main_layout.addWidget(video_panel, 7)

        # 右侧控制面板（原始代码完整保留）
        control_panel = QWidget()
        control_layout = QVBoxLayout()
        control_layout.setSpacing(20)
        control_layout.setContentsMargins(20, 20, 20, 20)

        # 输入框组（原始代码完整保留）
        input_group = QGroupBox("用户信息")
        input_layout = QFormLayout()
        self.name_edit = QLineEdit()
        self.name_edit.setPlaceholderText("输入姓名...")
        self.id_edit = QLineEdit()
        self.id_edit.setPlaceholderText("输入学号...")
        input_layout.addRow("姓名：", self.name_edit)
        input_layout.addRow("学号：", self.id_edit)
        input_group.setLayout(input_layout)
        control_layout.addWidget(input_group)

        # 控制按钮组（原始代码完整保留）
        btn_group = QWidget()
        btn_layout = QHBoxLayout()
        self.play_pause_btn = QPushButton("开始识别")
        self.play_pause_btn.setIcon(QIcon.fromTheme("media-playback-start"))
        self.play_pause_btn.clicked.connect(self.toggle_play_pause)
        self.capture_btn = QPushButton("采集人脸")
        self.capture_btn.setIcon(QIcon.fromTheme("camera-photo"))
        self.capture_btn.clicked.connect(self.capture_face)
        self.punch_btn = QPushButton("立即打卡")
        self.punch_btn.setIcon(QIcon.fromTheme("system-run"))
        self.punch_btn.clicked.connect(self.punch_in)
        # 添加训练按钮
        train_btn = QPushButton("训练模型")
        train_btn.clicked.connect(self.train_model)
        btn_layout.addWidget(self.play_pause_btn)
        btn_layout.addWidget(self.capture_btn)
        btn_layout.addWidget(self.punch_btn)
        btn_layout.addWidget(train_btn)
        btn_group.setLayout(btn_layout)
        control_layout.addWidget(btn_group)

        # 状态显示区（原始代码完整保留，包含 status_label）
        status_group = QGroupBox("系统状态")
        status_layout = QVBoxLayout()
        self.status_label = QLabel("就绪 | 摄像头正常")  # 原始定义保留
        self.status_label.setStyleSheet("color: white; font-size: 16px;")
        self.conf_slider = QSlider(Qt.Horizontal)
        self.conf_slider.setRange(60, 95)
        self.conf_slider.setValue(self.conf_threshold)
        self.conf_slider.valueChanged.connect(self.update_threshold)
        status_layout.addWidget(self.status_label)
        status_layout.addWidget(QLabel("识别灵敏度："))
        status_layout.addWidget(self.conf_slider)
        status_group.setLayout(status_layout)
        control_layout.addWidget(status_group)

        # 学生信息管理部分
        manage_group = QGroupBox("学生信息管理")
        manage_layout = QVBoxLayout()
        self.student_list = QListWidget()

        # 新增一个QWidget来容纳按钮布局
        btn_widget = QWidget()
        btn_hbox = QHBoxLayout(btn_widget)  # 将布局关联到btn_widget
        delete_btn = QPushButton("删除学生")
        delete_btn.clicked.connect(self.delete_student)
        modify_btn = QPushButton("修改信息")
        modify_btn.clicked.connect(self.modify_student)
        btn_hbox.addWidget(delete_btn)
        btn_hbox.addWidget(modify_btn)

        manage_layout.addWidget(self.student_list)
        manage_layout.addWidget(btn_widget)  # 添加QWidget而不是布局
        manage_group.setLayout(manage_layout)
        control_layout.addWidget(manage_group)

        control_panel.setLayout(control_layout)
        main_layout.addWidget(control_panel, 3)

        self.setLayout(main_layout)

    def update_frame(self):
        """原始视频刷新逻辑完整保留"""
        if self.is_paused:
            return

        ret, frame = self.cap.read()
        if not ret:
            self.status_label.setText("错误：无法访问摄像头！")
            self.status_label.setStyleSheet("color: #e74c3c;")
            return

        # 人脸检测与绘制（原始代码完整保留）
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)

        for (x, y, w, h) in faces:
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
            if self.play_pause_btn.text() == "停止识别":
                face_roi = gray[y:y + h, x:x + w]
                try:
                    label, conf = self.face_recognizer.predict(face_roi)
                    if conf < self.conf_threshold:
                        user_id = self.label_map.get(label, "未知")
                        cv2.putText(frame, f"{user_id} ({100 - conf:.1f}%)",
                                    (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
                except:
                    pass

        # 转换图像格式（原始代码完整保留）
        h, w = frame.shape[:2]
        q_img = QImage(frame.data, w, h, 3 * w, QImage.Format_BGR888)
        self.video_label.setPixmap(QPixmap.fromImage(q_img).scaled(
            self.video_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))

    def toggle_play_pause(self):
        """原始播放/暂停逻辑完整保留"""
        self.is_paused = not self.is_paused
        if self.is_paused:
            self.play_pause_btn.setText("开始识别")
            self.play_pause_btn.setIcon(QIcon.fromTheme("media-playback-start"))
        else:
            self.play_pause_btn.setText("停止识别")
            self.play_pause_btn.setIcon(QIcon.fromTheme("media-playback-stop"))

    def update_threshold(self, value):
        """原始灵敏度更新逻辑完整保留"""
        self.conf_threshold = value
        self.status_label.setText(f"识别灵敏度：{self.conf_threshold}% | 状态：就绪")

    def load_model(self):
        """加载模型"""
        if os.path.exists("face_model.yml") and os.path.exists("label_map.pkl"):
            self.face_recognizer.read("face_model.yml")
            with open("label_map.pkl", "rb") as f:
                label_map = pickle.load(f)
            self.label_map = {v: k for k, v in label_map.items()}
            self.status_label.setText("模型加载成功")
            self.load_dataset()  # 加载标签映射和学生列表
        else:
            QMessageBox.warning(self, "缺少模型", "请先训练模型", QMessageBox.Ok)

    def load_dataset(self):
        """加载标签映射和学生列表"""
        label_map = {}
        if os.path.exists("label_map.pkl"):
            with open("label_map.pkl", "rb") as f:
                label_map = pickle.load(f)
        self.label_map = {v: k for k, v in label_map.items()}
        self.student_list.clear()
        for student_id in self.label_map.values():
            self.student_list.addItem(student_id)

    def capture_face(self):
        """原始采集逻辑（新增模型更新，最小改动）"""
        name = self.name_edit.text().strip()
        id_num = self.id_edit.text().strip()

        if not id_num:
            QMessageBox.warning(self, "输入错误", "学号不能为空！", QMessageBox.Ok)
            return

        ret, frame = self.cap.read()
        if not ret:
            QMessageBox.critical(self, "摄像头错误", "无法读取摄像头数据。", QMessageBox.Ok)
            return

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)

        if len(faces) == 0:
            QMessageBox.warning(self, "未检测到人脸", "请确保正对摄像头。", QMessageBox.Ok)
            return

        success_count = 0
        for i, (x, y, w, h) in enumerate(faces):
            face_img = gray[y:y + h, x:x + w]
            face_img = cv2.resize(face_img, (100, 100))
            timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
            file_name = f"pic/{id_num}-{timestamp}.jpg"
            cv2.imwrite(file_name, face_img)
            success_count += 1

        self.load_dataset()  # 更新学生列表
        QMessageBox.information(self, "采集成功",
                                f"成功保存 {success_count} 张人脸图片！",
                                QMessageBox.Ok)

    def punch_in(self):
        """原始打卡逻辑完整保留"""
        if self.is_paused:
            QMessageBox.warning(self, "系统暂停", "请先开始识别。", QMessageBox.Ok)
            return

        ret, frame = self.cap.read()
        if not ret:
            QMessageBox.critical(self, "摄像头错误", "无法读取摄像头数据。", QMessageBox.Ok)
            return

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)

        if len(faces) == 0:
            QMessageBox.warning(self, "未检测到人脸", "请确保正对摄像头。", QMessageBox.Ok)
            return

        for (x, y, w, h) in faces:
            face_roi = gray[y:y + h, x:x + w]
            try:
                label, conf = self.face_recognizer.predict(face_roi)
                if conf < self.conf_threshold:
                    user_id = self.label_map.get(label, "未知")
                    if user_id == "未知":
                        raise ValueError("未注册用户")
                    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    self.save_punch_record(user_id, current_time)

                    face_roi = cv2.resize(face_roi, (100, 100))  # 统一尺寸
                    timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
                    file_path = f"pic/{user_id}-{timestamp}.jpg"  # 用学号作为文件名前缀
                    cv2.imwrite(file_path, face_roi)
                    self.load_dataset()  # 更新学生列表

                    os.makedirs("attendance_images", exist_ok=True)  # 创建打卡图像目录
                    img_filename = f"attendance_images/{user_id}_{current_time.replace(':', '-')}.jpg"
                    cv2.putText(frame, f"打卡时间：{current_time}", (10, 30),  # 添加时间水印
                                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
                    cv2.imwrite(img_filename, frame)  # 保存带识别框的彩色图像

                    QMessageBox.information(self, "打卡成功",
                                            f"{user_id} 于 {current_time} 打卡成功！\n"
                                            f"图像已保存：{img_filename}",  # 提示保存路径
                                            QMessageBox.Ok)

                    QMessageBox.information(self, "打卡成功",
                                            f"{user_id} 于 {current_time} 打卡成功！",
                                            QMessageBox.Ok)
                    return
                else:
                    raise ValueError("识别失败")
            except:
                QMessageBox.warning(self, "识别失败", "未找到匹配的用户，请重新尝试。", QMessageBox.Ok)
                return

    def save_punch_record(self, user_id, time_str):
        """原始记录保存逻辑完整保留"""
        with open("attendance_records.csv", "a", encoding="utf-8-sig") as f:
            f.write(f"{user_id},{time_str}\n")

    def timerEvent(self, event):
        """原始定时器逻辑完整保留"""
        self.update_frame()

    def closeEvent(self, event):
        """原始关闭逻辑完整保留"""
        self.cap.release()
        event.accept()

    def delete_student(self):
        selected_items = self.student_list.selectedItems()
        if not selected_items:
            QMessageBox.warning(self, "未选择", "请选择要删除的学生", QMessageBox.Ok)
            return
        student_id = selected_items[0].text()
        pic_dir = 'pic'
        for filename in os.listdir(pic_dir):
            if filename.startswith(student_id + '-'):
                os.remove(os.path.join(pic_dir, filename))
        self.load_dataset()
        QMessageBox.information(self, "删除成功", f"已删除学生 {student_id} 的所有信息", QMessageBox.Ok)

    def modify_student(self):
        selected_items = self.student_list.selectedItems()
        if not selected_items:
            QMessageBox.warning(self, "未选择", "请选择要修改的学生", QMessageBox.Ok)
            return
        old_id = selected_items[0].text()
        dialog = QDialog(self)
        dialog.setWindowTitle("修改学生信息")
        layout = QFormLayout(dialog)
        new_name_edit = QLineEdit()
        new_name_edit.setPlaceholderText("新姓名")
        new_id_edit = QLineEdit()
        new_id_edit.setPlaceholderText("新学号")
        layout.addRow("新姓名：", new_name_edit)
        layout.addRow("新学号：", new_id_edit)
        button_box = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        layout.addWidget(button_box)
        button_box.accepted.connect(dialog.accept)
        button_box.rejected.connect(dialog.reject)
        if dialog.exec_() == QDialog.Accepted:
            new_name = new_name_edit.text().strip()
            new_id = new_id_edit.text().strip()
            if not new_id:
                QMessageBox.warning(dialog, "输入错误", "新学号不能为空", QMessageBox.Ok)
                return
            pic_dir = 'pic'
            for filename in os.listdir(pic_dir):
                if filename.startswith(old_id + '-'):
                    old_path = os.path.join(pic_dir, filename)
                    new_filename = filename.replace(old_id, new_id)
                    new_path = os.path.join(pic_dir, new_filename)
                    os.rename(old_path, new_path)
            self.load_dataset()
            QMessageBox.information(self, "修改成功", f"已将学生 {old_id} 的信息修改为 {new_id}", QMessageBox.Ok)

    def train_model(self):
        """训练模型"""
        images = []
        labels = []
        label_map = {}
        current_label = 0

        if not os.path.exists('pic'):
            os.makedirs('pic')

        for filename in os.listdir('pic'):
            if filename.endswith('.jpg'):
                id_num = filename.split('-')[0]
                if id_num not in label_map:
                    label_map[id_num] = current_label
                    current_label += 1
                label = label_map[id_num]

                img_path = os.path.join('pic', filename)
                img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
                if img is not None:
                    img = cv2.resize(img, (100, 100))
                    images.append(img)
                    labels.append(label)

        if len(images) > 0:
            self.face_recognizer.train(images, np.array(labels))
            self.face_recognizer.save("face_model.yml")
            with open("label_map.pkl", "wb") as f:
                pickle.dump(label_map, f)
            self.label_map = {v: k for k, v in label_map.items()}
            self.load_dataset()  # 更新标签映射和学生列表
            QMessageBox.information(self, "训练完成", "模型训练成功", QMessageBox.Ok)
        else:
            QMessageBox.warning(self, "无数据", "没有可训练的人脸数据", QMessageBox.Ok)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    app.setStyle("Fusion")
    ex = FaceRecognitionSystem()
    ex.show()
    sys.exit(app.exec_())